import argparse
import json
import os
import sys
from base64 import b64encode
from io import BytesIO
from pathlib import Path

import uvicorn
from fastapi import FastAPI, responses
from starlette.responses import JSONResponse

now_dir = os.getcwd()
sys.path.append(now_dir)
from dotenv import load_dotenv
from scipy.io import wavfile

from configs.config import Config
from infer.modules.vc.modules import VC

####
# USAGE
#
# In your Terminal or CMD or whatever


def arg_parse() -> tuple:
    parser = argparse.ArgumentParser()
    parser.add_argument("--f0up_key", type=int, default=0)
    parser.add_argument("--input_path", type=str, help="input path")
    parser.add_argument("--index_path", type=str, help="index path")
    parser.add_argument("--f0method", type=str, default="harvest", help="harvest or pm")
    parser.add_argument("--opt_path", type=str, help="opt path")
    parser.add_argument("--model_name", type=str, help="store in assets/weight_root")
    parser.add_argument("--index_rate", type=float, default=0.66, help="index rate")
    parser.add_argument("--device", type=str, help="device")
    parser.add_argument("--is_half", type=bool, help="use half -> True")
    parser.add_argument("--filter_radius", type=int, default=3, help="filter radius")
    parser.add_argument("--resample_sr", type=int, default=0, help="resample sr")
    parser.add_argument("--rms_mix_rate", type=float, default=1, help="rms mix rate")
    parser.add_argument("--protect", type=float, default=0.33, help="protect")
    args = parser.parse_args()
    sys.argv = sys.argv[:1]
    return args

load_dotenv()
app = FastAPI()

@app.get("/voiceAsync1")
def inference1(
    input_audio: str,
    modelpath: str,
    res_type: str = None,
    sid: int = 0,
    f0_up_key: int = 0,
    f0_method: str = 'rmvpe',
    f0_file: str = None,
    index_file: str = None,
    index_rate: float = 0.75,
    filter_radius: int = 3,
    resample_sr: int = 0,
    rms_mix_rate: float = 0.25,
    protect: float = 0.33,
):
    print("==================="+index_file)
    print(res_type)
    config = Config()
    vc = VC(config)
    vc.get_vc1(Path(modelpath))
    tgt_sr, audio_opt, times, _ = vc.vc_inference(
        sid,
        input_audio,
        f0_up_key,
        f0_method,
        f0_file,
        index_file,
        index_rate,
        filter_radius,
        resample_sr,
        rms_mix_rate,
        protect,
    )
    wavfile.write(wv := BytesIO(), tgt_sr, audio_opt)
    print(times)
    if res_type == "blob":
        return responses.StreamingResponse(
            wv,
            media_type="audio/wav",
            headers={"Content-Disposition": "attachment; filename=inference.wav"},
        )
    else:
        return JSONResponse(
            {
                "time": json.loads(json.dumps(times)),
                "audio": b64encode(wv.read()).decode("utf-8"),
            }
        )

@app.get("/voiceAsync")
def inference(
    input_audio: str,
    modelpath: str,
    res_type: str = None,
    sid: int = 0,
    f0_up_key: int = 0,
    f0_method: str = 'rmvpe',
    f0_file: str = None,
    index_file: str = None,
    index_rate: float = 0.75,
    filter_radius: int = 3,
    resample_sr: int = 0,
    rms_mix_rate: float = 0.25,
    protect: float = 0.33,
):
    print("==================="+index_file)
    print(res_type)
    config = Config()
    config.device = "cuda:0"
    vc = VC(config)
    vc.get_vc1(Path(modelpath))
    _, wav_opt = vc.vc_single(
        0,
        input_audio,
        f0_up_key,
        f0_file,
        f0_method,
        index_file,
        None,
        index_rate,
        filter_radius,
        resample_sr,
        rms_mix_rate,
        protect,
    )
    wavfile.write(wv := BytesIO(), wav_opt[0], wav_opt[1])
    if res_type == "blob":
        return responses.StreamingResponse(
            wv,
            media_type="audio/wav",
            headers={"Content-Disposition": "attachment; filename=inference.wav"},
        )
    else:
        return JSONResponse(
            {
                "audio": b64encode(wv.read()).decode("utf-8"),
            }
        )

def main():
    load_dotenv()
    args = arg_parse()
    config = Config()
    config.device = args.device if args.device else config.device
    config.is_half = args.is_half if args.is_half else config.is_half
    vc = VC(config)
    vc.get_vc(args.model_name)
    _, wav_opt = vc.vc_single(
        0,
        args.input_path,
        args.f0up_key,
        None,
        args.f0method,
        args.index_path,
        None,
        args.index_rate,
        args.filter_radius,
        args.resample_sr,
        args.rms_mix_rate,
        args.protect,
    )
    wavfile.write(args.opt_path, wav_opt[0], wav_opt[1])


if __name__ == "__main__":
    os.environ["rmvpe_root"]="assets/rmvpe"
    print(os.environ["rmvpe_root"])
    uvicorn.run(app, host='127.0.0.1', port=8001, workers=1)
